
import os
import random
import shutil
import cv2
from networkx import triangular_lattice_graph
from tqdm import tqdm


def get_data_list():
    save_path = '/media/yw/SDA3/nnUnet_dataset/origin_data_1031/main_segment_data/dataset'
    use_file = '/media/yw/SDA3/nnUnet_dataset/origin_data_1031/main_segment_data/train/visual'
    use_list =  os.listdir(use_file)
    origin_file = '/media/yw/SDA3/nnUnet_dataset/origin_data_1031/main_segment_data/train/images'
    LAD_file = '/media/yw/SDA3/nnUnet_dataset/origin_data_1031/main_segment_data/train/LAD_labels'
    LCX_file = '/media/yw/SDA3/nnUnet_dataset/origin_data_1031/main_segment_data/train/LCX_labels'
    train_list, val_list = split_dataset(use_list)
    for i in tqdm(train_list):
        origin = os.path.join(origin_file,i)
        LAD = os.path.join(LAD_file,i)
        LCX = os.path.join(LCX_file,i)

        images = os.path.join(save_path,'images','train',i)
        LAD_MASK = os.path.join(save_path,'LAD_mask','train',i)
        LCX_MASK = os.path.join(save_path,'LCX_mask','train',i)
        os.makedirs(os.path.join(save_path,'images','train'), exist_ok=True)
        os.makedirs(os.path.join(save_path,'LAD_mask','train'), exist_ok=True)
        os.makedirs(os.path.join(save_path,'LCX_mask','train'), exist_ok=True)

        shutil.copyfile(origin,images)
        LAD = cv2.imread(LAD,0)
        LCX = cv2.imread(LCX,0)
        LAD[LAD > 0] = 1
        LCX[LCX > 0] = 1
        cv2.imwrite(LAD_MASK,LAD)
        cv2.imwrite(LCX_MASK,LCX)
    for i in tqdm(val_list):
        origin = os.path.join(origin_file,i)
        LAD = os.path.join(LAD_file,i)
        LCX = os.path.join(LCX_file,i)

        images = os.path.join(save_path,'images','val',i)
        LAD_MASK = os.path.join(save_path,'LAD_mask','val',i)
        LCX_MASK = os.path.join(save_path,'LCX_mask','val',i)
        os.makedirs(os.path.join(save_path,'images','val'), exist_ok=True)
        os.makedirs(os.path.join(save_path,'LAD_mask','val'), exist_ok=True)
        os.makedirs(os.path.join(save_path,'LCX_mask','val'), exist_ok=True)

        shutil.copyfile(origin,images)
        LAD = cv2.imread(LAD,0)
        LCX = cv2.imread(LCX,0)
        LAD[LAD > 0] = 1
        LCX[LCX > 0] = 1
        cv2.imwrite(LAD_MASK,LAD)
        cv2.imwrite(LCX_MASK,LCX)

def split_dataset(use_list):
    num = len(use_list)
    train_num = int(num * 0.8)
    val_num = int(num * 0.2)
    random.seed(77)
    use_list = random.sample(use_list,num)
    train_list = use_list[:train_num]
    val_list = use_list[train_num:]
    return train_list, val_list


def normliaztion_seg():
    import cv2
    paths = [
        # "/media/yw/SDA2/zhongnan/dataset/main_vessel_segment/main_vessel/dataset/LAD_mask/training",
        # "/media/yw/SDA2/zhongnan/dataset/main_vessel_segment/main_vessel/dataset/LAD_mask/validation",
        "/media/yw/SDA2/zhongnan/dataset/main_vessel_segment/main_vessel/dataset/LCX_mask/training",
        "/media/yw/SDA2/zhongnan/dataset/main_vessel_segment/main_vessel/dataset/LCX_mask/validation",
    ]
    for img_dir in paths:
        for img in os.listdir(img_dir):
            img_path = os.path.join(img_dir,img)
            img = cv2.imread(img_path,0)
            print(img)
            img = img / 255
            cv2.imwrite(img_path,img)



if __name__ == "__main__":
    get_data_list()
    # normliaztion_seg()